# -*- coding: UTF-8 -*-
"""
@Project ：swift-forge-engine 
@File    ：snow_flake
@IDE     ：PyCharm 
@Author  ：Tony
@Date    ：2025/2/19 14:24 
@Des     ：
"""
import time


class Snowflake:
    def __init__(self, worker_id, datacenter_id):
        # 自定义机器ID和数据中心ID，取值范围根据需要设定
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id

        # 定义位数
        self.timestamp_bits = 41
        self.worker_bits = 10
        self.datacenter_bits = 10
        self.sequence_bits = 12

        # 最大值
        self.max_worker_id = -1 ^ (-1 << self.worker_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_bits)

        # 上次时间戳
        self.last_timestamp = -1
        # 当前序列号
        self.sequence = 0

        # 时间戳左移位数
        self.timestamp_shift = self.worker_bits + self.datacenter_bits + self.sequence_bits
        # 机器ID左移位数
        self.worker_shift = self.datacenter_bits + self.sequence_bits
        # 数据中心ID左移位数
        self.datacenter_shift = self.sequence_bits

        # 序列号掩码
        self.sequence_mask = -1 ^ (-1 << self.sequence_bits)

    def _current_timestamp(self):
        return int(time.time() * 1000)

    def _wait_next_timestamp(self, last_timestamp):
        timestamp = self._current_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._current_timestamp()
        return timestamp

    def generate_id(self):
        timestamp = self._current_timestamp()

        if timestamp < self.last_timestamp:
            raise ValueError("Clock moved backwards.")

        if timestamp == self.last_timestamp:
            self.sequence = (self.sequence + 1) & self.sequence_mask
            if self.sequence == 0:
                timestamp = self._wait_next_timestamp(self.last_timestamp)
        else:
            self.sequence = 0

        self.last_timestamp = timestamp

        unique_id = ((timestamp << self.timestamp_shift) |
                     (self.datacenter_id << self.datacenter_shift) |
                     (self.worker_id << self.worker_shift) |
                     self.sequence)
        return unique_id


snowflake = Snowflake(worker_id=1, datacenter_id=1)


def generate_id() -> str:
    unique_id = snowflake.generate_id()
    return str(unique_id)